#pragma once

#include <vector>
#include <queue>
#include <iostream>
#include <deque>
#include <memory>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <future>
#include <functional>
#include <stdexcept>
#include <atomic>
#include <type_traits>
#include <random>
#include <algorithm>
#include <cassert>
#include <pthread.h>

/**
 * @class WorkStealingQueue
 * @brief 支持任务窃取的线程安全队列
 * 
 * 该队列实现了线程本地任务队列，允许其他线程"窃取"任务，
 * 通过减少对全局任务队列的访问来提高性能。
 */
class WorkStealingQueue 
{
public:
    /**
     * @brief 构造函数
     */
    WorkStealingQueue() = default;
    
    /**
     * @brief 析构函数
     */
    ~WorkStealingQueue() = default;
    
    /**
     * @brief 禁止拷贝构造
     */
    WorkStealingQueue(const WorkStealingQueue&) = delete;
    
    /**
     * @brief 禁止赋值操作
     */
    WorkStealingQueue& operator=(const WorkStealingQueue&) = delete;
    
    /**
     * @brief 向队列尾部添加任务
     * @param task 要添加的任务
     */
    void push(std::function<void()> task) 
    {
        std::lock_guard<std::mutex> lock(mutex_);
        tasks_.push_back(std::move(task));
    }
    
    /**
     * @brief 从队列头部获取任务（供本地线程使用）
     * @param task 用于存储获取的任务
     * @return 是否成功获取任务
     */
    bool pop(std::function<void()>& task) 
    {
        std::lock_guard<std::mutex> lock(mutex_);
        if (tasks_.empty()) 
        {
            return false;
        }
        
        task = std::move(tasks_.back());
        tasks_.pop_back();
        return true;
    }
    
    /**
     * @brief 从队列尾部窃取任务（供其他线程使用）
     * @param task 用于存储窃取的任务
     * @return 是否成功窃取任务
     */
    bool steal(std::function<void()>& task) 
    {
        std::lock_guard<std::mutex> lock(mutex_);
        if (tasks_.empty()) 
        {
            return false;
        }
        
        task = std::move(tasks_.front());
        tasks_.pop_front();
        return true;
    }
    
    /**
     * @brief 检查队列是否为空
     * @return 队列是否为空
     */
    bool empty() const 
    {
        std::lock_guard<std::mutex> lock(mutex_);
        return tasks_.empty();
    }
    
    /**
     * @brief 获取队列中任务数量
     * @return 任务数量
     */
    size_t size() const 
    {
        std::lock_guard<std::mutex> lock(mutex_);
        return tasks_.size();
    }

private:
    // 任务队列
    std::deque<std::function<void()>> tasks_;
    // 队列互斥锁
    mutable std::mutex mutex_;
};

/**
 * @class StdThreadPool
 * @brief 高性能线程池实现
 */
class StdThreadPool {
public:
    /**
     * @brief 任务优先级枚举
     */
    enum class TaskPriority 
    {
        LOW,    // 低优先级
        NORMAL, // 普通优先级
        HIGH    // 高优先级（嵌套任务使用）
    };

    /**
     * @brief 构造函数，创建指定数量的工作线程
     * @param threads 线程数量，默认为硬件支持的并发线程数
     */
    explicit StdThreadPool(size_t threads = std::thread::hardware_concurrency());
    
    /**
     * @brief 向线程池提交任务
     * @tparam F 函数类型
     * @tparam Args 参数类型包
     * @param f 要执行的函数
     * @param args 函数参数
     * @param priority 任务优先级，默认为普通优先级
     * @return std::future<返回值类型> 用于获取任务执行结果
     */
    template<class F, class... Args>
    auto enqueue(F&& f, Args&&... args, TaskPriority priority = TaskPriority::NORMAL) 
        -> std::future<typename std::invoke_result<F, Args...>::type>;
    
    /**
     * @brief 获取当前等待执行的任务数量
     * @return 任务数量
     */
    size_t get_pending_tasks() const;
    
    /**
     * @brief 获取线程池中的线程数量
     * @return 线程数量
     */
    size_t get_thread_count() const;
    
    /**
     * @brief 析构函数，等待所有任务完成并销毁线程
     */
    ~StdThreadPool();
    
    /**
     * @brief 禁止拷贝构造
     */
    StdThreadPool(const StdThreadPool&) = delete;
    
    /**
     * @brief 禁止赋值操作
     */
    StdThreadPool& operator=(const StdThreadPool&) = delete;

private:
    /**
     * @brief 任务结构体
     */
    struct Task 
    {
        std::function<void()> func;
        TaskPriority priority;
        
        Task(std::function<void()> f, TaskPriority p = TaskPriority::NORMAL)
            : func(std::move(f)), priority(p) {}
            
        // 优先级比较，用于优先级队列
        bool operator<(const Task& other) const 
        {
            return priority < other.priority;
        }
    };

    /**
     * @brief 工作线程函数
     * @param worker_id 工作线程ID
     */
    void worker_thread(size_t worker_id);
    
    /**
     * @brief 尝试从全局队列获取任务
     * @param task 用于存储获取的任务
     * @return 是否成功获取任务
     */
    bool get_task_from_global_queue(std::function<void()>& task);
    
    /**
     * @brief 尝试从高优先级队列获取任务
     * @param task 用于存储获取的任务
     * @return 是否成功获取任务
     */
    bool get_task_from_priority_queue(std::function<void()>& task);
    
    /**
     * @brief 尝试从其他线程窃取任务
     * @param worker_id 当前工作线程ID
     * @param task 用于存储窃取的任务
     * @return 是否成功窃取任务
     */
    bool steal_task_from_other_thread(size_t worker_id, std::function<void()>& task);
    
    /**
     * @brief 提交带优先级的任务到适当的队列
     * @param task 任务函数
     * @param priority 任务优先级
     */
    void submit_task(std::function<void()> task, TaskPriority priority);

    // 工作线程数组
    std::vector<std::thread> workers_;
    
    // 线程本地任务队列数组
    std::vector<std::unique_ptr<WorkStealingQueue>> local_work_queues_;
    
    // 全局任务队列
    std::queue<std::function<void()>> global_tasks_;
    
    // 高优先级任务队列
    std::deque<std::function<void()>> priority_tasks_;
    
    // 同步原语
    mutable std::mutex global_queue_mutex_;
    mutable std::mutex priority_queue_mutex_;
    std::condition_variable condition_;
    
    // 线程池状态标志
    std::atomic<bool> stop_;
    
    // 当前等待执行的任务数量
    mutable std::atomic<size_t> pending_tasks_;
    
    // 随机数生成器，用于任务窃取时随机选择线程
    mutable std::mt19937 rng_;

    // 线程ID映射
    std::vector<std::thread::id> thread_ids_;
    std::mutex thread_ids_mutex_;
    
    // 嵌套深度跟踪，用于防止死锁
    static thread_local int task_nesting_depth_;
    
    // 避免嵌套调用中的自锁和死锁
    static thread_local bool in_task_execution_;
    
    // 每个线程的ID缓存
    static thread_local size_t thread_local_id_;
};

// 初始化线程本地静态变量
thread_local bool StdThreadPool::in_task_execution_ = false;
thread_local int StdThreadPool::task_nesting_depth_ = 0;
thread_local size_t StdThreadPool::thread_local_id_ = static_cast<size_t>(-1);

/**
 * @brief 线程池构造函数实现
 * @param threads 线程数量
 */
inline StdThreadPool::StdThreadPool(size_t threads)
    : stop_(false), pending_tasks_(0), rng_(std::random_device{}())
{
    size_t num_threads = (threads > 0) ? threads : std::thread::hardware_concurrency();
    
    // 至少创建一个线程
    num_threads = std::max(size_t(1), num_threads);
    
    // 初始化线程本地任务队列
    local_work_queues_.resize(num_threads);
    for(size_t i = 0; i < num_threads; ++i) {
        local_work_queues_[i] = std::make_unique<WorkStealingQueue>();
    }
    
    // 预留空间以避免重新分配
    workers_.reserve(num_threads);
    thread_ids_.resize(num_threads);
    
    try 
    {
        // 创建工作线程
        for(size_t i = 0; i < num_threads; ++i) 
        {
            workers_.emplace_back(&StdThreadPool::worker_thread, this, i);
        }
    } 
    catch(...) 
    {
        // 如果创建线程失败，清理并重新抛出异常
        stop_ = true;
        condition_.notify_all();
        for(auto& worker : workers_) 
        {
            if(worker.joinable()) 
            {
                worker.join();
            }
        }
        throw;
    }

    pthread_atfork(nullptr, nullptr, +[]{
        task_nesting_depth_ = 0;
        in_task_execution_  = false;
        thread_local_id_    = size_t(-1);
    });
}

/**
 * @brief 从全局队列获取任务
 */
inline bool StdThreadPool::get_task_from_global_queue(std::function<void()>& task) 
{
    std::lock_guard<std::mutex> lock(global_queue_mutex_);
    if(global_tasks_.empty()) 
    {
        return false;
    }
    
    task = std::move(global_tasks_.front());
    global_tasks_.pop();
    return true;
}

/**
 * @brief 从高优先级队列获取任务
 */
inline bool StdThreadPool::get_task_from_priority_queue(std::function<void()>& task) 
{
    std::lock_guard<std::mutex> lock(priority_queue_mutex_);
    if(priority_tasks_.empty()) 
    {
        return false;
    }
    
    task = std::move(priority_tasks_.front());
    priority_tasks_.pop_front();
    return true;
}

/**
 * @brief 从其他线程窃取任务
 */
inline bool StdThreadPool::steal_task_from_other_thread(size_t worker_id, std::function<void()>& task) 
{
    // 生成随机索引，避免总是从同一线程窃取任务
    std::uniform_int_distribution<size_t> dist(0, local_work_queues_.size() - 1);
    
    // 尝试从其他线程窃取任务
    size_t num_queues = local_work_queues_.size();
    for(size_t i = 0; i < num_queues; ++i) 
    {
        size_t index = (dist(rng_) + i) % num_queues;
        if(index == worker_id || !local_work_queues_[index]) 
        {
            continue;
        }
        
        if(local_work_queues_[index]->steal(task)) 
        {
            return true;
        }
    }
    
    return false;
}

/**
 * @brief 工作线程函数实现
 * @param worker_id 工作线程ID
 */
inline void StdThreadPool::worker_thread(size_t worker_id) 
{
    // 存储线程ID以便识别
    {
        std::lock_guard<std::mutex> lock(thread_ids_mutex_);
        thread_ids_[worker_id] = std::this_thread::get_id();
        thread_local_id_ = worker_id;
    }
    
    // 获取线程本地任务队列
    auto& local_queue = local_work_queues_[worker_id];
    
    while(!stop_.load(std::memory_order_relaxed) || pending_tasks_.load(std::memory_order_relaxed) > 0) 
    {
        std::function<void()> task;
        bool got_task = false;
        
        // 1. 首先尝试从高优先级队列获取任务
        if(get_task_from_priority_queue(task)) 
        {
            got_task = true;
        }
        // 2. 如果高优先级队列为空，尝试从本地队列获取任务
        else if(local_queue && local_queue->pop(task)) 
        {
            got_task = true;
        } 
        // 3. 如果本地队列为空，尝试从全局队列获取任务
        else if(get_task_from_global_queue(task)) 
        {
            got_task = true;
        } 
        // 4. 如果全局队列也为空，尝试从其他线程窃取任务
        else if(steal_task_from_other_thread(worker_id, task))
        {
            got_task = true;
        } 
        // 5. 如果所有队列都为空，等待新任务或终止信号
        else 
        {
            std::unique_lock<std::mutex> lock(global_queue_mutex_);
            
            // 如果已经收到停止信号且没有待处理任务，则退出
            if(stop_.load(std::memory_order_relaxed) && pending_tasks_.load(std::memory_order_relaxed) == 0) 
            {
                break;
            }
            
            // 设置超时，避免线程永久等待
            auto timeout = std::chrono::milliseconds(100);
            condition_.wait_for(lock, timeout, [this, &task, &got_task]{ 
                // 先检查高优先级队列
                {
                    std::lock_guard<std::mutex> priority_lock(priority_queue_mutex_);
                    if(!priority_tasks_.empty()) 
                    {
                        task = std::move(priority_tasks_.front());
                        priority_tasks_.pop_front();
                        got_task = true;
                        return true;
                    }
                }
                
                // 再检查全局队列
                if(!global_tasks_.empty()) 
                {
                    task = std::move(global_tasks_.front());
                    global_tasks_.pop();
                    got_task = true;
                    return true;
                }
                
                // 如果线程池停止且无任务，则退出等待
                return stop_.load(std::memory_order_relaxed) && pending_tasks_.load(std::memory_order_relaxed) == 0;
            });
            
            // 如果超时，继续循环
            if(!got_task) 
            {
                continue;
            }
        }
        
        // 执行任务
        if(got_task) 
        {
            pending_tasks_.fetch_sub(1, std::memory_order_relaxed);
            
            // 防止任务执行时的异常
            in_task_execution_ = true;
            ++task_nesting_depth_; // 增加嵌套深度
            
            try 
            {
                task();
            } 
            catch(const std::exception& e) 
            {
                // 记录异常但不终止线程
                std::cerr << "任务执行异常: " << e.what() << std::endl;
            } 
            catch(...) 
            {
                // 防止任务异常导致线程终止
                std::cerr << "任务执行发生未知异常" << std::endl;
            }
            
            --task_nesting_depth_; // 减少嵌套深度
            in_task_execution_ = false;
        }
    }
}

/**
 * @brief 提交任务到适当的队列
 */
inline void StdThreadPool::submit_task(std::function<void()> task, TaskPriority priority) 
{
    // 获取当前线程ID
    size_t worker_id = thread_local_id_;
    bool is_worker_thread = (worker_id != static_cast<size_t>(-1));
    
    // 增加待处理任务计数
    pending_tasks_.fetch_add(1, std::memory_order_relaxed);
    
    // 根据优先级和线程状态决定放入哪个队列
    if(priority == TaskPriority::HIGH) 
    {
        // 高优先级任务放入专门的优先级队列
        std::lock_guard<std::mutex> lock(priority_queue_mutex_);
        priority_tasks_.push_back(std::move(task));
        condition_.notify_one();
    }
    else if(!in_task_execution_ && is_worker_thread && 
            worker_id < local_work_queues_.size() && local_work_queues_[worker_id]) 
    {
        // 非嵌套任务且是工作线程，放入本地队列
        local_work_queues_[worker_id]->push(std::move(task));
        condition_.notify_one();
    }
    else 
    {
        // 嵌套任务或非工作线程，放入全局队列
        std::unique_lock<std::mutex> lock(global_queue_mutex_);
        
        // 不允许在线程池停止后添加任务
        if(stop_.load(std::memory_order_relaxed)) 
        {
            pending_tasks_.fetch_sub(1, std::memory_order_relaxed);
            lock.unlock();
            throw std::runtime_error("enqueue on stopped StdThreadPool");
        }
        
        // 添加任务到全局队列
        global_tasks_.emplace(std::move(task));
        lock.unlock();
        condition_.notify_one();
    }
}

/**
 * @brief 向线程池提交任务
 */
template<class F, class... Args>
auto StdThreadPool::enqueue(F&& f, Args&&... args, TaskPriority priority) 
    -> std::future<typename std::invoke_result<F, Args...>::type>
{
    // 定义返回类型
    using return_type = typename std::invoke_result<F, Args...>::type;
    
    // 创建任务包装
    auto task = std::make_shared<std::packaged_task<return_type()>>(
        std::bind(std::forward<F>(f), std::forward<Args>(args)...)
    );
    
    // 获取future对象
    std::future<return_type> result = task->get_future();
    
    // 嵌套任务使用高优先级，避免死锁
    if(task_nesting_depth_ > 0) 
    {
        priority = TaskPriority::HIGH;
    }
    
    // 包装任务并提交
    submit_task([task](){ (*task)(); }, priority);
    
    return result;
}

/**
 * @brief 获取等待执行的任务数量
 */
inline size_t StdThreadPool::get_pending_tasks() const 
{
    return pending_tasks_.load(std::memory_order_relaxed);
}

/**
 * @brief 获取线程池中的线程数量
 */
inline size_t StdThreadPool::get_thread_count() const 
{
    return workers_.size();
}

/**
 * @brief 线程池析构函数
 */
inline StdThreadPool::~StdThreadPool() 
{
    {
        std::unique_lock<std::mutex> lock(global_queue_mutex_);
        stop_.store(true, std::memory_order_relaxed);
    }
    
    // 通知所有等待的线程
    condition_.notify_all();
    
    // 等待所有线程完成
    for(std::thread &worker : workers_) 
    {
        if(worker.joinable()) 
        {
            worker.join();
        }
    }
}